{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3D LIDAR Object Detection using TAO PointPillars\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of PointPillars\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/sample_2.png?raw=true_2.png\" width=\"720\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Learning Objectives\n",
    "In this notebook, you learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a PointPillars model on the KITTI dataset\n",
    "* Prune the trained model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Run evaluation & inference on the trained model to verify the accuracy\n",
    "* Export & deploy the model in TensorRT\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `PointPillars` model, \n",
    "which you may deploy with TensorRT. For more information on how to deploy this model, visit this\n",
    "[GitHub](https://github.com/NVIDIA-AI-IOT/viz_3Dbbox_ros2_pointpillars) repository sample.\n",
    "\n",
    " ### Table of Contents\n",
    "\n",
    " This notebook shows an example use case of PointPillars using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    " 0. [Set up env variables and map drives](#head-0)\n",
    " 1. [Install the TAO launcher](#head-1)\n",
    " 2. [Prepare dataset and pretrained model](#head-2)<br>\n",
    "     2.1 [Download the dataset](#head-2-1)<br>\n",
    "     2.2 [Verify the downloaded dataset](#head-2-2)<br>\n",
    "     2.3 [Convert dataset to required format](#head-2-3)<br>\n",
    " 3. [Provide training specification](#head-3)\n",
    " 4. [Run TAO training](#head-4)\n",
    " 5. [Evaluate trained models](#head-5)\n",
    " 6. [Prune trained models](#head-6)\n",
    " 7. [Retrain pruned models](#head-7)\n",
    " 8. [Evaluate retrained model](#head-8)\n",
    " 9. [Visualize inferences](#head-9)\n",
    " 10. [Deploy](#head-10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    " \n",
    "The following notebook requires you to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the your workspace. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "print(\"Please replace the variables with your own.\")\n",
    "%env KEY=tlt_encode\n",
    "\n",
    "# Define this local project directory that must be mapped to the TAO Docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/tao-experiments\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data/pointpillars\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"pointpillars\"\n",
    ")\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/pointpillars\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data/pointpillars\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/pointpillars\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/pointpillars/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO Docker instance, so that the data and the results are mapped from in and out of the Docker. For more information, see the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO Docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    },\n",
    "    # set gpu index for tao-converter\n",
    "    \"Envs\": [\n",
    "        {\"variable\": \"CUDA_VISIBLE_DEVICES\", \"value\": os.getenv(\"GPU_INDEX\")},\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO Docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a Python package distributed as a Python wheel listed in PyPI. You install the launcher by executing the following cell.\n",
    "\n",
    "The TAO Toolkit recommend that you run the TAO launcher in a virtual env with Python 3.6.9. You can follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a Python virtual env using the `virtualenv` and `virtualenvwrapper` packages. After you have setup virtualenvwrapper, set the version of Python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO Python package, you must meet the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 460+\n",
    "\n",
    "Once you have installed the pre-requisites, log in to the Docker registry nvcr.io by following the command below:\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You are prompted to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Skip this step if you have already installed the TAO launcher.\n",
    "!pip3 install --upgrade nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 2. Prepare dataset and pretrained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the KITTI detection dataset for the tutorial. To find more details, see\n",
    " http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Download the KITTI detection images (http://www.cvlibs.net/download.php?file=data_object_image_2.zip), labels(http://www.cvlibs.net/download.php?file=data_object_label_2.zip), velodyne LIDAR pointcloud(http://www.cvlibs.net/download.php?file=data_object_velodyne.zip), and LIDAR calibration file(http://www.cvlibs.net/download.php?file=data_object_calib.zip) to $DATA_DOWNLOAD_DIR.\n",
    " \n",
    "The data is extracted with the following structure:\n",
    "\n",
    "```bash\n",
    "│── ImageSets\n",
    "│── training\n",
    "│   ├──calib & velodyne & label_2 & image_2\n",
    "│── testing\n",
    "│   ├──calib & velodyne & image_2\n",
    "```\n",
    "\n",
    "The `testing` directory is not be used in this notebook because it has no labels. For the `training` dataset, a script is used to do data preprocessing and split it into `train` and `val` splits. Finally, verify that the directory seen by TAO PointPillars looks like:\n",
    "\n",
    "```bash\n",
    "│── train\n",
    "│   ├──lidar & label\n",
    "│── val\n",
    "│   ├──lidar & label\n",
    "```\n",
    "\n",
    "You can also use this notebook with your own dataset. To use this example with your own dataset, use the same directory structure as mentioned below."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have gotten the download links in your email, populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL`,  `KITTI_LABELS_DOWNLOAD_URL`, `KITTI_LIDAR_DOWNLOAD_DIR`, and `KITTI_CALIB_DOWNLOAD_DIR`. This next cell, downloads the data and place in `$LOCAL_DATA_DIR`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that images are only required for KITTI dataset in this notebook, but not required for a general dataset that follows TAO PointPillars standard format. The reason that we need images in the KITTI dataset is that the KITTI dataset does not conform with the standard format and some pre-processing is necessary for it. The preprocessing reads each image's size and retrieves only points that are in the field-of-view(FOV) of the camera from the original LiDAR files. The retrieved FOV-only points are saved to a new LiDAR file for each of the original LiDAR files. This is necessary because the KITTI dataset only has labels in the FOV of camera, but no labels for points outside of camera FOV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "\n",
    "os.environ[\"URL_IMAGES\"]=\"KITTI_IMAGES_DOWNLOAD_URL\"\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $LOCAL_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "\n",
    "os.environ[\"URL_LABELS\"]=\"KITTI_LABELS_DOWNLOAD_URL\"\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $LOCAL_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi\n",
    "\n",
    "os.environ[\"URL_LIDAR\"]=\"KITTI_LIDAR_DOWNLOAD_URL\"\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_velodyne.zip ]; then wget $URL_LIDAR -O $LOCAL_DATA_DIR/data_object_velodyne.zip; else echo \"velodyne archive already downloaded\"; fi \n",
    "\n",
    "os.environ[\"URL_CALIB\"]=\"CALIB_DOWNLOAD_URL\"\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_calib.zip ]; then wget $URL_CALIB -O $LOCAL_DATA_DIR/data_object_calib.zip; else echo \"calib archive already downloaded\"; fi "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_velodyne.zip ]; then echo 'Velodyne zip file not found, please download.'; else echo 'Found Velodyne zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_calib.zip ]; then echo 'Calib zip file not found, please download.'; else echo 'Found Calib zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_image_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_label_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_velodyne.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_calib.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Convert dataset to required format<a class=\"anchor\" id=\"head-2-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output directories\n",
    "!mkdir -p $LOCAL_DATA_DIR/train/lidar\n",
    "!mkdir -p $LOCAL_DATA_DIR/train/label\n",
    "!mkdir -p $LOCAL_DATA_DIR/val/lidar\n",
    "!mkdir -p $LOCAL_DATA_DIR/val/label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Retrieve FOV-only LIDAR points from 360-degree LIDAR points\n",
    "# Since only FOV data is labelled in KITTI dataset\n",
    "!tao model pointpillars run python $SPECS_DIR/gen_lidar_points.py -p $DATA_DOWNLOAD_DIR/training/velodyne \\\n",
    "                                                           -c $DATA_DOWNLOAD_DIR/training/calib    \\\n",
    "                                                           -i $DATA_DOWNLOAD_DIR/training/image_2  \\\n",
    "                                                           -o $DATA_DOWNLOAD_DIR/train/lidar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert labels from Camera coordinate system to LIDAR coordinate system, etc\n",
    "!tao model pointpillars run python $SPECS_DIR/gen_lidar_labels.py -l $DATA_DOWNLOAD_DIR/training/label_2 \\\n",
    "                                                           -c $DATA_DOWNLOAD_DIR/training/calib \\\n",
    "                                                           -o $DATA_DOWNLOAD_DIR/train/label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop DontCare class\n",
    "!tao model pointpillars run python $SPECS_DIR/drop_class.py $DATA_DOWNLOAD_DIR/train/label DontCare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train/val split\n",
    "!tao model pointpillars run python $SPECS_DIR/kitti_split.py $SPECS_DIR/val.txt \\\n",
    "                                                      $DATA_DOWNLOAD_DIR/train/lidar \\\n",
    "                                                      $DATA_DOWNLOAD_DIR/train/label \\\n",
    "                                                      $DATA_DOWNLOAD_DIR/val/lidar \\\n",
    "                                                      $DATA_DOWNLOAD_DIR/val/label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate dataset statistics for data augmentation\n",
    "!tao model pointpillars dataset_convert -e $SPECS_DIR/pointpillars.yaml results_dir=$USER_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/pointpillars.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    " * Provide the sample spec file for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model pointpillars train -e $SPECS_DIR/pointpillars.yaml \\\n",
    "                              results_dir=$USER_EXPERIMENT_DIR \\\n",
    "                              dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                              key=$KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to run multi-gpu training\n",
    "# !tao model pointpillars train -e $SPECS_DIR/pointpillars.yaml \\\n",
    "#                               results_dir=$USER_EXPERIMENT_DIR \\\n",
    "#                               dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "#                               key=$KEY \\\n",
    "#                               train.num_gpus=2 \\"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The evaluation metric in TAO PointPillars is different from that in official metric of KITTI point cloud detection. While KITTI metric considers easy/moderate/hard categories of objects and filters small objects whose sizes are smaller than a threshold, it is only meaningful for KITTI dataset. Instead, TAO PointPillars metric is a general metric that does not classify objects into easy/moderate/hard categories and does not exclude objects in calculation of metric. This makes TAO PointPillars metric a general metric that is applicable to a general dataset. The final result is average precision (AP) and mean average precision (mAP) regardless of its details in computation. Due to this, the TAO PointPillars metric is not comparable with KITTI official metric on KITTI dataset, although they are typically similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model pointpillars evaluate -e $SPECS_DIR/pointpillars.yaml \\\n",
    "                                 results_dir=$USER_EXPERIMENT_DIR \\\n",
    "                                 dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                                 key=$KEY"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    " * Specify pre-trained model\n",
    " * Threshold for pruning\n",
    " * A key to save and load the model\n",
    " * Output directory to store the model\n",
    " \n",
    "Usually, you must to adjust `prune.pruning_thresh` (threshold) for accuracy and model size trade off. Using a higher threshold gives you smaller model (and thus higher inference speed) but accuracy suffers. The threshold to use depends on the dataset. The threshold value below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model pointpillars prune -e $SPECS_DIR/pointpillars.yaml \\\n",
    "                              results_dir=$USER_EXPERIMENT_DIR \\\n",
    "                              dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                              key=$KEY \\\n",
    "                              prune.pruning_thresh=0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -lht $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    " * Model needs to be re-trained to bring back accuracy after pruning\n",
    " * Specify re-training specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model pointpillars train -e $SPECS_DIR/pointpillars_retrain.yaml \\\n",
    "                              results_dir=$USER_EXPERIMENT_DIR/retrain \\\n",
    "                              dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                              key=$KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -lht $LOCAL_EXPERIMENT_DIR/retrain"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model pointpillars evaluate -e $SPECS_DIR/pointpillars_retrain.yaml \\\n",
    "                                 results_dir=$USER_EXPERIMENT_DIR/retrain \\\n",
    "                                 dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                                 key=$KEY"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    " In this section, we run the inference command on the trained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model pointpillars inference -e $SPECS_DIR/pointpillars_retrain.yaml \\\n",
    "                                  results_dir=$USER_EXPERIMENT_DIR/retrain \\\n",
    "                                  dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                                  key=$KEY \\\n",
    "                                  inference.save_to_file=True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `inference` command produces visualization of bounding boxes of objects in and rendering of point cloud. This command can be slow due to plots. If you are not going to finish it, you can abort it and check the partially detected results(images) in `$USER_EXPERIMENT_DIR/retrain/detected_boxes`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 10. Deploy! <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Export in FP16 mode.\n",
    "# Note that the .onnx model in FP16 mode is the same as in FP32 mode.\n",
    "!tao model pointpillars run rm -f $USER_EXPERIMENT_DIR/retrain/checkpoint_epoch_80.onnx\n",
    "!tao model pointpillars export -e $SPECS_DIR/pointpillars_retrain.yaml \\\n",
    "                        dataset.data_info_path=$USER_EXPERIMENT_DIR/data_info \\\n",
    "                        key=$KEY \\\n",
    "                        results_dir=$USER_EXPERIMENT_DIR/retrain \\\n",
    "                        export.save_engine=$USER_EXPERIMENT_DIR/retrain/checkpoint_epoch_80.trt"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython2",
  "version": 2,
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
